import pymysql
import plotly.express as px

import pandas as pd

db_config={
        'host': '192.168.31.134',
        'user': 'root',
        'password': 'Password123@mysql',
        'database': 'oecd_data'
    }


connection = pymysql.connect(
    host=db_config.get('host', 'localhost'),
    user=db_config.get('user', 'root'),
    password=db_config.get('password', ''),
    database=db_config.get('database', 'test'),
    charset='utf8mb4'
)

cursor = connection.cursor()

per_capita_consumptio_sql = '''
SELECT 
  country_name, time_period, obs_value 
FROM 
  `gdp_consumption_per_capita`
WHERE
  sector = 'S14'
AND
  transaction_2 = 'Final consumption expenditure, per capita'
'''

work_hour_sql = '''
SELECT country_name, time_period, obs_value FROM `annual_hours_worked_per_year`
WHERE worker_status = '_T'
'''

df_per_capita_consumption = pd.read_sql(per_capita_consumptio_sql, connection)
print('>>> 顺利加载per_capita_consumption数据')
df_work_hour = pd.read_sql(work_hour_sql, connection)
print('>>> 顺利加载work_hour数据')

connection.close()


# 方法2：使用交集
common_keys = set(zip(df_per_capita_consumption['country_name'], 
                      df_per_capita_consumption['time_period'])) & \
              set(zip(df_work_hour['country_name'], 
                      df_work_hour['time_period']))

# 过滤出共同的行
df_per_capita_consumption_common = df_per_capita_consumption[df_per_capita_consumption.apply(lambda row: (row['country_name'], row['time_period']) in common_keys, axis=1)]
print('>>> 顺利加载per_capita_consumption_common数据')
df_work_hour_common = df_work_hour[df_work_hour.apply(lambda row: (row['country_name'], row['time_period']) in common_keys, axis=1)]
print('>>> 顺利加载work_hour_common数据')
# 基于country_name和time_period进行合并
df_merged = df_per_capita_consumption_common.merge(df_work_hour_common, 
                               on=['country_name', 'time_period'], 
                               suffixes=('_per_capita_consumption', '_work_hour'))

df_out = df_merged[['country_name', 'time_period', 'obs_value_per_capita_consumption', 'obs_value_work_hour']]

duplicates = df_out.duplicated(subset=['country_name', 'time_period']).sum()
print(f"重复数据行数: {duplicates}")

df_scatter = df_out

# 确保时间周期是数值类型
df_scatter['time_period'] = df_scatter['time_period'].astype(int)


df_scatter['week_hour'] = df_scatter['obs_value_work_hour'] / 52
print(f"散点图数据行数: {len(df_scatter)}")
print(f"包含国家数: {df_scatter['country_name'].nunique()}")
print(f"时间范围: {df_scatter['time_period'].min()} - {df_scatter['time_period'].max()}")

df_scatter.to_excel('工作时长与人均消费.xlsx', index=False)

# 方法3：使用自定义图例设置
import plotly.graph_objects as go

# 获取所有国家
countries = df_scatter['country_name'].unique()

# 高级设置：添加更美观的参考线
# 获取y轴的范围
y_min = df_scatter['week_hour'].min()
y_max = df_scatter['week_hour'].max()

# 创建新的图表
fig_with_reference = go.Figure()

# 为每个国家添加数据
for country in countries:
    country_data = df_scatter[df_scatter['country_name'] == country]
    
    fig_with_reference.add_trace(go.Scatter(
        x=country_data['obs_value_per_capita_consumption'],
        y=country_data['week_hour'],
        mode='markers',
        name=country,
        marker=dict(size=8),
        text=country_data['time_period'].astype(str),
        hovertemplate='<b>%{customdata[0]}</b><br>年份: %{text}<br>人均消费: %{x:.2f}<br>周工作时长: %{y:.2f}<extra></extra>',
        customdata=country_data[['country_name']].values
    ))

# 添加竖着的虚线参考线
fig_with_reference.add_trace(go.Scatter(
    x=[10000, 10000],
    y=[y_min, y_max],
    mode='lines',
    line=dict(color='red', width=3, dash='dash'),
    name='人均消费=10000',
    showlegend=True,
    hovertemplate='人均消费参考线: 10000<extra></extra>'
))

# 在参考线上添加文本标签
fig_with_reference.add_annotation(
    x=10000,
    y=y_max,
    text="人均消费=10000",
    showarrow=True,
    arrowhead=2,
    ax=0,
    ay=-40,
    bgcolor="white",
    bordercolor="red",
    borderwidth=1
)

# 更新布局
fig_with_reference.update_layout(
    title='工作时长与人均消费关系（含参考线）',
    xaxis_title='人均消费',
    yaxis_title='周工作时长（小时）',
    legend_title='国家/参考线',
    showlegend=True,
    xaxis=dict(
        showgrid=True,
        gridwidth=1,
        gridcolor='LightGray'
    ),
    yaxis=dict(
        showgrid=True,
        gridwidth=1,
        gridcolor='LightGray'
    )
)

# fig_with_reference.show()
fig_with_reference.write_html('工作时长与人均消费散点图_高级参考线.html')
print("高级参考线散点图已保存为 '工作时长与人均消费散点图_高级参考线.html'")
